#include "associator.h"
#include "../tracker/tracker_manager.h"
#include "../utils/km.h"
#include "associate_factor.h"

#include <eigen3/Eigen/Core>

#include <iostream>
#include <set>
#include <unordered_map>

namespace fusion_perception {

Associator::Associator(std::shared_ptr<TrackerManager> trackerManager, std::vector<std::shared_ptr<AssociateFactor>> factors)
    : trackerManager_(trackerManager)
    , factors_(std::move(factors))
{
}

MultiAssociatedObject Associator::associate(int64_t microTs, const ObjectSptrVector& dest, const TimeSyncObject& src)
{
    auto destTmp = dest;
    std::set<ObjectSptr> newObjectSet;
    std::unordered_map<ObjectSptr, ObjectSptrVector> associatedObjectUMap;

    for (auto& i : src) {
        for (auto& j : associate(microTs, destTmp, i.second)) {
            if (j.first == nullptr) {
                j.first = j.second;
                destTmp.emplace_back(j.first);
                newObjectSet.emplace(j.first);
            }

            associatedObjectUMap[j.first].emplace_back(j.second);
        }
    }

    MultiAssociatedObject multiAssociatedObject;
    for (auto& i : associatedObjectUMap) {
        if (newObjectSet.find(i.first) != newObjectSet.end()) {
            multiAssociatedObject.emplace(nullptr, std::move(i.second));
        } else {
            multiAssociatedObject.emplace(i.first, std::move(i.second));
        }
    }

    return multiAssociatedObject;
}

AssociatedObject Associator::associate(int64_t, const ObjectSptrVector& dest, const ObjectSptrVector& src)
{
    AssociatedObject associatedObject;

    if (dest.empty()) {
        for (auto& i : src) {
            associatedObject.emplace_back(std::pair<ObjectSptr, ObjectSptr>(nullptr, i));
        }
        return associatedObject;
    }

    Eigen::MatrixXf associationMat(dest.size(), src.size());

    for (size_t i = 0; i < dest.size(); i++) {
        for (size_t j = 0; j < src.size(); j++) {
            associationMat(i, j) = assessScore(*dest[i], *src[j]);
        }
    }

    std::vector<std::pair<int, int>> matchedCps;

    if (associationMat.rows() > associationMat.cols()) {
        assert(KM().get(associationMat.transpose(), matchedCps, true));
    } else {
        assert(KM().get(associationMat, matchedCps, false));
    }

    for (auto& i : matchedCps) {
        if (i.second != -1) {
            if (i.first == -1) {
                associatedObject.emplace_back(std::pair<ObjectSptr, ObjectSptr>(nullptr, src[i.second]));
            } else {
                associatedObject.emplace_back(std::pair<ObjectSptr, ObjectSptr>(dest[i.first], src[i.second]));
            }
        }
    }

    return associatedObject;
}

int32_t Associator::assessScore(const Object& o1, const Object& o2)
{
    int32_t sumScore = 0;

    for (auto& i : factors_) {
        int32_t score = 0;
        auto tracker = trackerManager_->getTrackerById(o1.id);

        if (tracker == nullptr) {
            score = i->assessScore(o1, o2);
        } else {
            score = i->assessScore(tracker, o2);
        }

        if (score == -1) {
            return -1;
        }

        score = score * (double)i->getWeight() / 100;
        sumScore += score;
    }

    return sumScore;
}

}